// DHCP Library v0.3 - April 25, 2009
// Author: Jordan Terrell - blog.jordanterrell.com

#include <Arduino.h>
#include "Ethernet.h"
#include "Dhcp.h"
#include "utility/w5100.h"

int DhcpClass::beginWithDHCP(uint8_t *mac, unsigned long timeout, unsigned long responseTimeout)
{
  _dhcpLeaseTime=0;
  _dhcpT1=0;
  _dhcpT2=0;
  _timeout = timeout;
  _responseTimeout = responseTimeout;

  // zero out _dhcpMacAddr
  memset(_dhcpMacAddr, 0, 6);
  reset_DHCP_lease();

  memcpy((void*)_dhcpMacAddr, (void*)mac, 6);
  _dhcp_state = STATE_DHCP_START;
  return request_DHCP_lease();
}

void DhcpClass::reset_DHCP_lease()
{
  // zero out _dhcpSubnetMask, _dhcpGatewayIp, _dhcpLocalIp, _dhcpDhcpServerIp, _dhcpDnsServerIp
  memset(_dhcpLocalIp, 0, 20);
}

  //return:0 on error, 1 if request is sent and response is received
int DhcpClass::request_DHCP_lease()
{
  uint8_t messageType = 0;

  // Pick an initial transaction ID
  _dhcpTransactionId = random(1UL, 2000UL);
  _dhcpInitialTransactionId = _dhcpTransactionId;

  _dhcpUdpSocket.stop();
  if (_dhcpUdpSocket.begin(DHCP_CLIENT_PORT) == 0) {
    // Couldn't get a socket
    return 0;
  }

  presend_DHCP();

  int result = 0;

  unsigned long startTime = millis();

  while (_dhcp_state != STATE_DHCP_LEASED) {
    if (_dhcp_state == STATE_DHCP_START) {
      _dhcpTransactionId++;
      send_DHCP_MESSAGE(DHCP_DISCOVER, ((millis() - startTime) / 1000));
      _dhcp_state = STATE_DHCP_DISCOVER;
    } else if (_dhcp_state == STATE_DHCP_REREQUEST) {
      _dhcpTransactionId++;
      send_DHCP_MESSAGE(DHCP_REQUEST, ((millis() - startTime)/1000));
      _dhcp_state = STATE_DHCP_REQUEST;
    } else if (_dhcp_state == STATE_DHCP_DISCOVER) {
      uint32_t respId;
      messageType = parseDHCPResponse(_responseTimeout, respId);
      if (messageType == DHCP_OFFER) {
        // We'll use the transaction ID that the offer came with,
        // rather than the one we were up to
        _dhcpTransactionId = respId;
        send_DHCP_MESSAGE(DHCP_REQUEST, ((millis() - startTime) / 1000));
        _dhcp_state = STATE_DHCP_REQUEST;
      }
    } else if (_dhcp_state == STATE_DHCP_REQUEST) {
      uint32_t respId;
      messageType = parseDHCPResponse(_responseTimeout, respId);
      if (messageType == DHCP_ACK) {
        _dhcp_state = STATE_DHCP_LEASED;
        result = 1;
        //use default lease time if we didn't get it
        if (_dhcpLeaseTime == 0) {
          _dhcpLeaseTime = DEFAULT_LEASE;
        }
        // Calculate T1 & T2 if we didn't get it
        if (_dhcpT1 == 0) {
          // T1 should be 50% of _dhcpLeaseTime
          _dhcpT1 = _dhcpLeaseTime >> 1;
        }
        if (_dhcpT2 == 0) {
          // T2 should be 87.5% (7/8ths) of _dhcpLeaseTime
          _dhcpT2 = _dhcpLeaseTime - (_dhcpLeaseTime >> 3);
        }
        _renewInSec = _dhcpT1;
        _rebindInSec = _dhcpT2;
      } else if (messageType == DHCP_NAK) {
        _dhcp_state = STATE_DHCP_START;
      }
    }

    if (messageType == 255) {
      messageType = 0;
      _dhcp_state = STATE_DHCP_START;
    }

    if (result != 1 && ((millis() - startTime) > _timeout))
      break;
  }

  // We're done with the socket now
  _dhcpUdpSocket.stop();
  _dhcpTransactionId++;

  _lastCheckLeaseMillis = millis();
  return result;
}

void DhcpClass::presend_DHCP()
{
}

void DhcpClass::send_DHCP_MESSAGE(uint8_t messageType, uint16_t secondsElapsed)
{
  uint8_t buffer[32];
  memset(buffer, 0, 32);
  IPAddress dest_addr(255, 255, 255, 255); // Broadcast address

  if (_dhcpUdpSocket.beginPacket(dest_addr, DHCP_SERVER_PORT) == -1) {
    //Serial.printf("DHCP transmit error\n");
    // FIXME Need to return errors
    return;
  }

  buffer[0] = DHCP_BOOTREQUEST;   // op
  buffer[1] = DHCP_HTYPE10MB;     // htype
  buffer[2] = DHCP_HLENETHERNET;  // hlen
  buffer[3] = DHCP_HOPS;          // hops

  // xid
  unsigned long xid = htonl(_dhcpTransactionId);
  memcpy(buffer + 4, &(xid), 4);

  // 8, 9 - seconds elapsed
  buffer[8] = ((secondsElapsed & 0xff00) >> 8);
  buffer[9] = (secondsElapsed & 0x00ff);

  // flags
  unsigned short flags = htons(DHCP_FLAGSBROADCAST);
  memcpy(buffer + 10, &(flags), 2);

  // ciaddr: already zeroed
  // yiaddr: already zeroed
  // siaddr: already zeroed
  // giaddr: already zeroed

  //put data in W5100 transmit buffer
  _dhcpUdpSocket.write(buffer, 28);

  memset(buffer, 0, 32); // clear local buffer

  memcpy(buffer, _dhcpMacAddr, 6); // chaddr

  //put data in W5100 transmit buffer
  _dhcpUdpSocket.write(buffer, 16);

  memset(buffer, 0, 32); // clear local buffer

  // leave zeroed out for sname && file
  // put in W5100 transmit buffer x 6 (192 bytes)

  for(int i = 0; i < 6; i++) {
    _dhcpUdpSocket.write(buffer, 32);
  }

  // OPT - Magic Cookie
  buffer[0] = (uint8_t)((MAGIC_COOKIE >> 24)& 0xFF);
  buffer[1] = (uint8_t)((MAGIC_COOKIE >> 16)& 0xFF);
  buffer[2] = (uint8_t)((MAGIC_COOKIE >> 8)& 0xFF);
  buffer[3] = (uint8_t)(MAGIC_COOKIE& 0xFF);

  // OPT - message type
  buffer[4] = dhcpMessageType;
  buffer[5] = 0x01;
  buffer[6] = messageType; //DHCP_REQUEST;

  // OPT - client identifier
  buffer[7] = dhcpClientIdentifier;
  buffer[8] = 0x07;
  buffer[9] = 0x01;
  memcpy(buffer + 10, _dhcpMacAddr, 6);

  // OPT - host name
  buffer[16] = hostName;
  buffer[17] = strlen(HOST_NAME) + 6; // length of hostname + last 3 bytes of mac address
  strcpy((char*)&(buffer[18]), HOST_NAME);

  printByte((char*)&(buffer[24]), _dhcpMacAddr[3]);
  printByte((char*)&(buffer[26]), _dhcpMacAddr[4]);
  printByte((char*)&(buffer[28]), _dhcpMacAddr[5]);

  //put data in W5100 transmit buffer
  _dhcpUdpSocket.write(buffer, 30);

  if (messageType == DHCP_REQUEST) {
    buffer[0] = dhcpRequestedIPaddr;
    buffer[1] = 0x04;
    buffer[2] = _dhcpLocalIp[0];
    buffer[3] = _dhcpLocalIp[1];
    buffer[4] = _dhcpLocalIp[2];
    buffer[5] = _dhcpLocalIp[3];

    buffer[6] = dhcpServerIdentifier;
    buffer[7] = 0x04;
    buffer[8] = _dhcpDhcpServerIp[0];
    buffer[9] = _dhcpDhcpServerIp[1];
    buffer[10] = _dhcpDhcpServerIp[2];
    buffer[11] = _dhcpDhcpServerIp[3];

    //put data in W5100 transmit buffer
    _dhcpUdpSocket.write(buffer, 12);
  }

  buffer[0] = dhcpParamRequest;
  buffer[1] = 0x06;
  buffer[2] = subnetMask;
  buffer[3] = routersOnSubnet;
  buffer[4] = dns;
  buffer[5] = domainName;
  buffer[6] = dhcpT1value;
  buffer[7] = dhcpT2value;
  buffer[8] = endOption;

  //put data in W5100 transmit buffer
  _dhcpUdpSocket.write(buffer, 9);

  _dhcpUdpSocket.endPacket();
}

uint8_t DhcpClass::parseDHCPResponse(unsigned long responseTimeout, uint32_t& transactionId)
{
  uint8_t type = 0;
  uint8_t opt_len = 0;

  unsigned long startTime = millis();

  while (_dhcpUdpSocket.parsePacket() <= 0) {
    if ((millis() - startTime) > responseTimeout) {
      return 255;
    }
    delay(50);
  }
  // start reading in the packet
  RIP_MSG_FIXED fixedMsg;
  _dhcpUdpSocket.read((uint8_t*)&fixedMsg, sizeof(RIP_MSG_FIXED));

  if (fixedMsg.op == DHCP_BOOTREPLY && _dhcpUdpSocket.remotePort() == DHCP_SERVER_PORT) {
    transactionId = ntohl(fixedMsg.xid);
    if (memcmp(fixedMsg.chaddr, _dhcpMacAddr, 6) != 0 ||
      (transactionId < _dhcpInitialTransactionId) ||
      (transactionId > _dhcpTransactionId)) {
      // Need to read the rest of the packet here regardless
      _dhcpUdpSocket.flush(); // FIXME
      return 0;
    }

    memcpy(_dhcpLocalIp, fixedMsg.yiaddr, 4);

    // Skip to the option part
    _dhcpUdpSocket.read((uint8_t *)NULL, 240 - (int)sizeof(RIP_MSG_FIXED));

    while (_dhcpUdpSocket.available() > 0) {
      switch (_dhcpUdpSocket.read()) {
      case endOption :
        break;

      case padOption :
        break;

      case dhcpMessageType :
        opt_len = _dhcpUdpSocket.read();
        type = _dhcpUdpSocket.read();
        break;

      case subnetMask :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read(_dhcpSubnetMask, 4);
        break;

      case routersOnSubnet :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read(_dhcpGatewayIp, 4);
        _dhcpUdpSocket.read((uint8_t *)NULL, opt_len - 4);
        break;

      case dns :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read(_dhcpDnsServerIp, 4);
        _dhcpUdpSocket.read((uint8_t *)NULL, opt_len - 4);
        break;

      case dhcpServerIdentifier :
        opt_len = _dhcpUdpSocket.read();
        if ( IPAddress(_dhcpDhcpServerIp) == IPAddress((uint32_t)0) ||
          IPAddress(_dhcpDhcpServerIp) == _dhcpUdpSocket.remoteIP() ) {
          _dhcpUdpSocket.read(_dhcpDhcpServerIp, sizeof(_dhcpDhcpServerIp));
        } else {
          // Skip over the rest of this option
          _dhcpUdpSocket.read((uint8_t *)NULL, opt_len);
        }
        break;

      case dhcpT1value :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read((uint8_t*)&_dhcpT1, sizeof(_dhcpT1));
        _dhcpT1 = ntohl(_dhcpT1);
        break;

      case dhcpT2value :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read((uint8_t*)&_dhcpT2, sizeof(_dhcpT2));
        _dhcpT2 = ntohl(_dhcpT2);
        break;

      case dhcpIPaddrLeaseTime :
        opt_len = _dhcpUdpSocket.read();
        _dhcpUdpSocket.read((uint8_t*)&_dhcpLeaseTime, sizeof(_dhcpLeaseTime));
        _dhcpLeaseTime = ntohl(_dhcpLeaseTime);
        _renewInSec = _dhcpLeaseTime;
        break;

      default :
        opt_len = _dhcpUdpSocket.read();
        // Skip over the rest of this option
        _dhcpUdpSocket.read((uint8_t *)NULL, opt_len);
        break;
      }
    }
  }

  // Need to skip to end of the packet regardless here
  _dhcpUdpSocket.flush(); // FIXME

  return type;
}


/*
    returns:
    0/DHCP_CHECK_NONE: nothing happened
    1/DHCP_CHECK_RENEW_FAIL: renew failed
    2/DHCP_CHECK_RENEW_OK: renew success
    3/DHCP_CHECK_REBIND_FAIL: rebind fail
    4/DHCP_CHECK_REBIND_OK: rebind success
*/
int DhcpClass::checkLease()
{
  int rc = DHCP_CHECK_NONE;

  unsigned long now = millis();
  unsigned long elapsed = now - _lastCheckLeaseMillis;

  // if more then one sec passed, reduce the counters accordingly
  if (elapsed >= 1000) {
    // set the new timestamps
    _lastCheckLeaseMillis = now - (elapsed % 1000);
    elapsed = elapsed / 1000;

    // decrease the counters by elapsed seconds
    // we assume that the cycle time (elapsed) is fairly constant
    // if the remainder is less than cycle time * 2
    // do it early instead of late
    if (_renewInSec < elapsed * 2) {
      _renewInSec = 0;
    } else {
      _renewInSec -= elapsed;
    }
    if (_rebindInSec < elapsed * 2) {
      _rebindInSec = 0;
    } else {
      _rebindInSec -= elapsed;
    }
  }

  // if we have a lease but should renew, do it
  if (_renewInSec == 0 &&_dhcp_state == STATE_DHCP_LEASED) {
    _dhcp_state = STATE_DHCP_REREQUEST;
    rc = 1 + request_DHCP_lease();
  }

  // if we have a lease or is renewing but should bind, do it
  if (_rebindInSec == 0 && (_dhcp_state == STATE_DHCP_LEASED ||
    _dhcp_state == STATE_DHCP_START)) {
    // this should basically restart completely
    _dhcp_state = STATE_DHCP_START;
    reset_DHCP_lease();
    rc = 3 + request_DHCP_lease();
  }
  return rc;
}

IPAddress DhcpClass::getLocalIp()
{
  return IPAddress(_dhcpLocalIp);
}

IPAddress DhcpClass::getSubnetMask()
{
  return IPAddress(_dhcpSubnetMask);
}

IPAddress DhcpClass::getGatewayIp()
{
  return IPAddress(_dhcpGatewayIp);
}

IPAddress DhcpClass::getDhcpServerIp()
{
  return IPAddress(_dhcpDhcpServerIp);
}

IPAddress DhcpClass::getDnsServerIp()
{
  return IPAddress(_dhcpDnsServerIp);
}

void DhcpClass::printByte(char * buf, uint8_t n )
{
  char *str = &buf[1];
  buf[0]='0';
  do {
    unsigned long m = n;
    n /= 16;
    char c = m - 16 * n;
    *str-- = c < 10 ? c + '0' : c + 'A' - 10;
  } while(n);
}
